'''
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.


Copyright (c) 2009 bryan p. conrad
bconrad2@gmail.com

Created on: 7/26/2009

Version: 0.1

@author: bryan
'''
#===============================================================================
# Imports
#===============================================================================
import JT 
from QVTKRenderWindowInteractor import * 
import numpy as np
import sys

#from msvcrt import kbhit, getch

# To stop render window from non-responding
# not sure if this is necessary with QT
import win32gui    


class DRR:
    """
    A collection of components that can be used to create a digitally
    reconstructed radiograph (DRR).

    With a decent graphics card the rendering can be hardware accelerated and does
    not necessarily need to be done on-screen.  A WindowToImageFilter is used 
    to capture the contents of the render window and save the projection as a 2D image.
    """

    def __init__(self):
        """
        Initialize DRR Class.   
        """
        #Set up Moving Image 3D renderwindow, with interaction
        self.ren = JT.vtk.vtkRenderer()
        self.renWin = JT.vtk.vtkRenderWindow()
        self.renWin.AddRenderer(self.ren)
        self.iren = JT.vtk.vtkRenderWindowInteractor()
        self.iren.SetRenderWindow(self.renWin)
        self.ren.SetLayer(0)
        self.ren.SetBackground(1,1,1)
        
        # Set default render window size 
        # To be set in the registration class once the size of the fixed 
        # image is known
        #self.renWin.SetSize(512,512)
        self.xrayCamera = JT.XRayCamera()
        self.ren.SetActiveCamera(self.xrayCamera)

        #We will need to use a custom interaction scheme (based on JointTrack)
        self.style = JT.JTInteractorStyle(self.iren)
        self.iren.SetInteractorStyle(self.style)

        #Create filter to capture an image from the render window
        self._w2if = JT.vtk.vtkWindowToImageFilter()
        self._w2if.SetInput(self.renWin)

        #Create a VTK Image Flipper
        self._vtkFlipperX = JT.vtk.vtkImageFlip()
        self._vtkFlipperY = JT.vtk.vtkImageFlip()

        self._vtkFlipperX.SetFilteredAxis(0)
        self._vtkFlipperY.SetFilteredAxis(1)

        #Create a VTK cast image filter to convert to float and char
        self._floatCast = JT.vtk.vtkImageCast()
        self._uCharCast = JT.vtk.vtkImageCast()

        self._floatCast.SetOutputScalarTypeToFloat()
        self._floatCast.ClampOverflowOn()

        self._uCharCast.SetOutputScalarTypeToUnsignedChar()
        self._uCharCast.ClampOverflowOn()

        # Fixed image size
        self._imageSize = []

        # Result of Generating an image is a projected image (2D VTK image)
        self._projectedImage = []

        # Create components for a background image
        self.renWin.SetNumberOfLayers(2)

        self._imageMapper = JT.vtk.vtkImageMapper()
        self._imageActor = JT.vtk.vtkActor2D()
        self._imageRen = JT.vtk.vtkRenderer()

        self.renWin.AddRenderer(self._imageRen)
        self._imageActor.SetMapper(self._imageMapper)
        self._imageRen.AddActor2D(self._imageActor)
        self._imageRen.SetLayer(1)
        self._imageRen.SetInteractive(0)
        self._imageRen.SetBackground(1,1,1)
        self._backgroundImage = None

        self._pose = (0,0,0,0,0,0)
        self.volume_list = []
        self.volume = JT.Volume3()
        self.volume_path = None
        self.ren.AddVolume(self.volume)

        
    def AddVolume(self,volume):
        """
        Add a volume to the DRR Renderer
        """
        #print 'DRR.AddVolume is Deprecated.  Use SetVolume instead.'
        self.ren.AddVolume(volume)
        #self.SetVolume(volume)
        
    def RemoveAllVolumes(self):
        """
        Remove all JT volumes (and any other props) from the renderer.
        """
        self.ren.RemoveAllViewProps()
        self.volume_list = []
        
    def SetVolume(self,volume):
        """
        Set a JT volume in the renderer.
        """
        if isinstance(volume,JT.Volume):
            self.RemoveAllVolumes()
            self.ren.AddVolume(volume)
            self.volume = volume
        elif isinstance(volume,JT.SurfaceModel):
            print "Adding Surface model to renderer"
            self.RemoveAllVolumes()
            self.ren.AddActor(volume)
            self.volume = volume
        else:
            print "DRR.SetVolume is not sure what to do with type: %s" % str(type(volume)) 
            return
        
        if self.volume.volume_path:
            self.volume_path = self.volume.volume_path
        else:
            self.volume_path = None
    
    def SetVolumeFileName(self,volume_path):
        self.volume_path = volume_path
        self.SetFileName(volume_path)
    
    def LoadSTL(self,filename):
        
        if not isinstance(self.volume,JT.SurfaceModel):
            self.volume = None
            self.volume = JT.SurfaceModel()
        self.SetFileName(filename)
        self.ren.RemoveAllViewProps()
        self.ren.AddActor(self.volume)
        
    def SetFileName(self,volume_path):
        if self.volume:
            self.volume_path = volume_path
            self.volume.SetVolumeFileName(volume_path)
        else:
            print 'ERROR. No volume class is associated with this instance of DRR'
    
    def AddBodyCenter(self,center):
        print "Deprecated.  Use SetCenter instead."
        self.volume.SetCenter(center)
    
    def SetCenter(self,center):
        self.volume.SetCenter(center)
    
    def AddSphereAtBodyCenter(self,center,color=(1,0,0)):
        # Add axes to renderer
        sphere = JT.vtk.vtkSphereSource()
        sphere.SetRadius(5.0)
        sphere.SetPhiResolution(10)
        mapper = JT.vtk.vtkPolyDataMapper()
        sphere_actor = JT.vtk.vtkActor()
        mapper.SetInput(sphere.GetOutput())
        sphere_actor.SetMapper(mapper)
        sphere_actor.SetPosition(center)
        sphere_actor.SetUserTransform(self.volume.GetUserTransform())
        props = sphere_actor.GetProperty()
        props.SetColor(color)
        self.ren.AddActor(sphere_actor)
    
    def AddAxesAtModelOrigin(self,model):
        # Add axes to renderer
        axes_actor = AxesWidget(model)
        axes_actor.SetInteractor(self.iren)
        self.ren.AddActor(axes_actor)
        origin = model.GetOrigin()
        
    def SetRenderWindowSize(self,size):
        """
        size is a tuple of integers that represents the size of the render
        window in pixels. 
        """
        self.renWin.SetSize(size)
    
    def Render(self):
        """
        Display renderwindow.  Does not start interaction.
        TODO: need to check to make sure that everything is setup.
        """
        win32gui.PeekMessage(0,0,0,0)
        self.renWin.Render()
        
    def StartInteraction(self):
        """
        Start the render window interactor.  If used in a script, the user must 
        quit the interaction (press q, or e typically) to continue to the 
        next command.
        """
        self.iren.Start()
    
    def UseVTKInteractionStyle(self):
        """
        Warning: This style of interaction will change the calibration setup,
        By allowing the user to move the position of the camera.  Do not use this 
        function for image registration.
        
        This method will turn on the VTK interactor style trackball camera.
        """
        style = JT.vtk.vtkInteractorStyleTrackballCamera()
        self.iren.SetInteractorStyle(style)
        
        
    def GenerateImage(self):
        """
        Render the volume and capture the resulting image from the render window.
        A change in this version of jt is to NOT convert from VTK image to ITK image
        """
        #To avoid being recognized as a non-responding program in windows, 
        #we need to check in on the renderwindow at each iteration
        win32gui.PeekMessage(0,0,0,0)
        
        #Capture projection from render window
        self.renWin.Render()
        self._w2if.Modified()

        #for image to be displayed correctly, it has to be flipped about the X axis
        #self._vtkFlipperX.SetInput(self._w2if.GetOutput())
        #self._vtkFlipperY.SetInput(self._vtkFlipperX.GetOutput())

        #self._projectedImage = self._vtkFlipperY.GetOutput()
        self._projectedImage = self._w2if.GetOutput()
        self._projectedImage.Update()
        
        #self._pyProjectedImage = JT.utils.vtk2numpy(self._projectedImage)
        #self._pyProjectedImage = JT.utils.vtk2numpy(self._projectedImage)

    def SaveGeneratedImage(self,filename):
        """
        Save the generated image to disk.  If a projection does not exist,
        it will be generated and saved, otherwise this function will save the
        image currently residing in self._projectedImage.

        The appropriate vtk image file writer is used, as determined 
        by the filename extension.
        """
        #print "Generating Image..."
        self.GenerateImage()
        JT.utils.genericImageWriter(self._projectedImage, filename)
        
                        
    def SetBlackOnWhiteImage(self):
        """
        This function changes the backgound color of the renderer.
        The renderer seems to automatically change the color of the actor so that
        it shows up on the background.
        """
        self._imageRen.SetBackground(1,1,1)
        self.Render()
        self._blackOnWhiteFlag = 1

    def SetWhiteOnBlackImage(self):
        """
        This function changes the backgound color of the renderer.
        The renderer seems to automatically change the color of the actor so that
        it shows up on the background.

        This is the default behavior.
        """
        self._imageRen.SetBackground(0,0,0)
        self.Render()
        self._blackOnWhiteFlag = 0

#    def GetXRayCamera(self):
#        """
#        Return the active Camera for the scene
#        """
#        return self.xrayCamera
#
#    def SetXRayCamera(self,camera):
#        """
#        Set the active Camera for the scene
#        """
#        self.xrayCamera = camera
#        self.ren.SetActiveCamera(camera)
#        
#        #Determine whether an internal or external calibration setup is being used
#        try:
#            cal=self.xrayCamera.IntCal
#            cal.isValid()
#        except:
#            try:
#                cal=self.xrayCamera.ExtCal
#                cal.isValid()
#            except:
#                print "DRR.GenerateImage: Cannot Generate Image, No calibration set"
#                print "External Cal: %s, Internal Cal: %s" %\
#                               (self.xrayCamera.IntCal.isValid(),self.xrayCamera.ExtCal.isValid())
#    def GetInternalCalibration(self):
#        """
#        The internal Calibration currently being used for the active camera
#        """
#        return self.xrayCamera.IntCal
#
#    def GetExternalCalibration(self):
#        """
#        The external Calibration currently being used for the active camera
#        """
#        return self.xrayCamera.ExtCal
#    
#    def SetExternalCalibration(self,cal):
#        self.xrayCamera.SetExternalCalibration(cal)
    def GetImage(self):
        """
        Return the current projected image (VTK format),
        The python (numpy) version can be obtained by using self._pyProjectedImage
        """
        #Assume that the user wants a new image to be generated.
        self.GenerateImage()
        return self._projectedImage
        
    def GetRenderTime(self):
        """
        Get the last render time in seconds (from renderer)
        """
        return self.ren.GetLastRenderTimeInSeconds()
    
#    def GetInteractor(self):
#        return self.iren
#    
    def LoadCalFile(self,cal_path,cal_type=None):
        """
        Set an a calibration in the DRR.  cal_type does not have to be given
        the file will be opened and the type will be read from the first 
        word in the file, either JT_INTCALIB or JT_EXTCALIB. 
        """
        # open the supplied cal file and take a peek at the first word 
        # to see what type of cal file it is
        file = open(cal_path)
        raw_data = file.read()
        file.close()
        data = raw_data.split()
        if data[0] == 'JT_INTCALIB':
            try:
                cal = self.xrayCamera.IntCal
                cal.LoadCalFile(cal_path)
                cal.cal_type = 'internal'
                self.xrayCamera.SetInternalCalibration(cal)
                self.xrayCamera.DoCalibration()
            except AttributeError:
                print 'ERROR. A component is missing, Calibration could not be set.'
                raise
            
        elif data[0] == 'JT_EXTCALIB':
            try:
                cal = self.xrayCamera.ExtCal
                cal.LoadCalFile(cal_path)
                cal.cal_type = 'external'
                self.xrayCamera.SetExternalCalibration(cal)
                self.xrayCamera.DoCalibration()
            except AttributeError:
                print 'ERROR. A component is missing, Calibration could not be set.'
                raise
        else:
            print "%s - Not recognized as a valid internal or external"\
                  "JointTrack calibration file"
            
    
    def LoadStaFile(self,sta_path):
        """
        Set the sta file location for volume position details.  It seems like
        this should be rolled into the consolidated calibration method, or is
        that not possible since this data is associated with a volume and 
        the other cal data is associated with a fluoro image?
        
        The Cal file must be loaded first to create Cal object in 
        the xrayCamera class.
        """
        try:
            cal = self.xrayCamera.Cal
            cal.LoadStaFile(sta_path)
            self.volume.SetOriginalTransform(cal.volumeTransform)
        except AttributeError:
            print 'ERROR. A component is missing, Calibration could not be set.'
            raise
    
    
    def SetBackgroundImage(self,vtkImage,flip=False):
        """
        Supply a VKT image to be displayed as the background in the same render
        window with the DRR projection.
        """
        #if a background image is used, imageRen needs to be moved to the back layer
        if vtkImage is None:
            self._imageRen.RemoveActor(self._imageActor)
            # After removing imageActor, background turned Black???
            self._imageRen.SetBackground(1,1,1)
            self.Render()
            return
        self.ren.SetLayer(1)
        self._imageRen.SetLayer(0)
        self._imageRen.AddActor2D(self._imageActor)
        self._backgroundImage = vtkImage
        
        flipperX = JT.vtk.vtkImageFlip()
        flipperY = JT.vtk.vtkImageFlip()
        flipperX.SetFilteredAxis(0)
        flipperY.SetFilteredAxis(1)
        flipperX.SetInput(vtkImage)
        flipperY.SetInput(flipperX.GetOutput())
        flipperY.Update()
        if flip is True:
            self._imageMapper.SetInput(flipperY.GetOutput())
        else:
            self._imageMapper.SetInput(vtkImage)

        (minPixel, maxPixel) = self._backgroundImage.GetScalarRange()

        self._imageMapper.SetColorLevel(maxPixel/2)
        self._imageMapper.SetColorWindow(maxPixel)

        self.Render()

    def InteractionOff(self):
        """
        Turn interaction off, useful for optimizing registration.  After user
        has specified a starting position
        """
        self.renWin.SetInteractor(None)

    def InteractionOn(self):
        """
        Turn interaction on. For user to specify a starting position
        """
        self.iren.SetRenderWindow(self.renWin)
    
    def OffScreenRenderingOn(self):
        self.renWin.SetOffScreenRendering(1)
    
    def OffScreenRenderingOff(self):
        self.renWin.SetOffScreenRendering(0)
            
class QTDRR(DRR):
    """
    A modified DRR that can be used in QT GUI applications 
    """
    def __init__(self,parent_object):
        #DRR.__init__(self)
        self.iren = QVTKRenderWindowInteractor(parent_object)
        self.ren = JT.vtk.vtkRenderer()
        self.renWin = self.iren.GetRenderWindow()
        self.renWin.AddRenderer(self.ren)
        self.iren.SetRenderWindow(self.renWin)
        self.ren.SetLayer(0)
        self.ren.SetBackground(1,1,1)
        
        # Set default render window size 
        # To be set in the registration class once the size of the fixed 
        # image is known
        #self.renWin.SetSize(512,512)
        self.xrayCamera = JT.XRayCamera()
        self.ren.SetActiveCamera(self.xrayCamera)

        #We will need to use a custom interaction scheme (based on JointTrack)
        self.style = JT.InteractorStyle.QTJTInteractorStyle(self.iren)
        self.iren.SetInteractorStyle(self.style)

        #Create filter to capture an image from the render window
        self._w2if = JT.vtk.vtkWindowToImageFilter()
        self._w2if.SetInput(self.renWin)

        #Create a VTK Image Flipper
        self._vtkFlipperX = JT.vtk.vtkImageFlip()
        self._vtkFlipperY = JT.vtk.vtkImageFlip()

        self._vtkFlipperX.SetFilteredAxis(0)
        self._vtkFlipperY.SetFilteredAxis(1)

        #Create a VTK cast image filter to convert to float and char
        self._floatCast = JT.vtk.vtkImageCast()
        self._uCharCast = JT.vtk.vtkImageCast()

        self._floatCast.SetOutputScalarTypeToFloat()
        self._floatCast.ClampOverflowOn()

        self._uCharCast.SetOutputScalarTypeToUnsignedChar()
        self._uCharCast.ClampOverflowOn()

        # Fixed image size
        self._imageSize = []

        # Result of Generating an image is a projected image (2D VTK image)
        self._projectedImage = []

        # Create components for a background image
        self.renWin.SetNumberOfLayers(2)

        self._imageMapper = JT.vtk.vtkImageMapper()
        self._imageActor = JT.vtk.vtkActor2D()
        self._imageRen = JT.vtk.vtkRenderer()

        self.renWin.AddRenderer(self._imageRen)
        self._imageActor.SetMapper(self._imageMapper)
        self._imageRen.AddActor2D(self._imageActor)
        self._imageRen.SetLayer(1)
        self._imageRen.SetInteractive(0)
        self._imageRen.SetBackground(1,1,1)

        self._backgroundImage = None

        self._pose = (0,0,0,0,0,0)
        self.volume_list = []
        # ***This is a change from the original DRR class***
        # Volume2 is better suited to real world situations when you don;t
        # know the true position and orientation of the volume.
        self.volume = JT.Volume2()
        self.ren.AddVolume(self.volume)
        self.volume_path = None
        
        #Setup a playback widget
#        playbackRep = drrPlaybackRep()
#        self.playbackWidget = JT.vtk.vtkPlaybackWidget()
#        self.playbackWidget.SetRepresentation(playbackRep)
#        self.playbackWidget.SetInteractor(self.iren)
#        self.playbackWidget.On()

class AnimateDRR(DRR):
    def __init__(self):
        #Set up Moving Image 3D renderwindow, with interaction
        self.ren = JT.vtk.vtkRenderer()
        self.renWin = JT.vtk.vtkRenderWindow()
        self.renWin.AddRenderer(self.ren)
        self.iren = JT.vtk.vtkRenderWindowInteractor()
        self.iren.SetRenderWindow(self.renWin)
        self.ren.SetLayer(0)
        #self.ren.SetBackground(1,1,1)
        
        #Create filter to capture an image from the render window
        self._w2if = JT.vtk.vtkWindowToImageFilter()
        self._w2if.SetInput(self.renWin)
        
        # Set default render window size 
        # To be set in the registration class once the size of the fixed 
        # image is known
        #self.renWin.SetSize(512,512)
        self.xrayCamera = JT.XRayCamera()
        self.ren.SetActiveCamera(self.xrayCamera)

        # Fixed image size
        self._imageSize = []

        # Result of Generating an image is a projected image (2D VTK image)
        self._projectedImage = []

        # Create components for a background image
        self.renWin.SetNumberOfLayers(2)

        self._imageMapper = JT.vtk.vtkImageMapper()
        self._imageActor = JT.vtk.vtkActor2D()
        self._imageRen = JT.vtk.vtkRenderer()

        self.renWin.AddRenderer(self._imageRen)
        self._imageActor.SetMapper(self._imageMapper)
        self._imageRen.AddActor2D(self._imageActor)
        self._imageRen.SetLayer(1)
        self._imageRen.SetInteractive(0)
        self._imageRen.SetBackground(0,0,0)
        self._backgroundImage = None
        
class QTAnimateDRR(AnimateDRR):
    def __init__(self,parent):
        self.iren = QVTKRenderWindowInteractor(parent)
        AnimateDRR.__init__(self)

class AxesActor(vtk.vtkAssembly):
    def __init__(self,length=50):
        """
        the axes will be associated with the model that is given.
        """
        Xaxis = self.make_arrow((1,0,0),'X', length, (1,0,0))
        Yaxis = self.make_arrow((0,1,0),'Y', length, (0,1,0))
        Zaxis = self.make_arrow((0,0,1),'Z', length, (0,0,1))
        self.AddPart(Xaxis)
        self.AddPart(Yaxis)
        self.AddPart(Zaxis)

    def ToggleVisibility(self):
        if self.GetVisibility():
            self.VisibilityOff()
        else:
            self.VisibilityOn()
    
    def make_arrow(self,axis,label,length=10,color=(1,1,1)):
        arrow_assembly = vtk.vtkAssembly()
        axis = np.asarray(axis)
        
        line_source = vtk.vtkLineSource()
        line_source.SetPoint1(0,0,0)
        line_source.SetPoint2(axis*length)
        line_mapper=vtk.vtkPolyDataMapper()
        line_mapper.SetInput(line_source.GetOutput())
        line_actor=vtk.vtkActor()
        line_actor.SetMapper(line_mapper)
        line_props = line_actor.GetProperty()
        line_props.SetColor(color)
        
        tip_source=vtk.vtkConeSource()
        tip_source.SetHeight(10)
        tip_source.SetRadius(5)
        tip_mapper=vtk.vtkPolyDataMapper()
        tip_mapper.SetInput(tip_source.GetOutput())
        tip_actor=vtk.vtkActor()
        tip_actor.SetMapper(tip_mapper)
        tip_actor.SetPosition(axis)
        tip_actor.RotateWXYZ(90,axis[0],axis[1],axis[2])
        tip_props = line_actor.GetProperty()    
        tip_props.SetColor(color)
        
        label_actor=vtk.vtkTextActor3D()
        label_actor.SetInput(label)
        label_actor.SetPosition((axis*length))
        
        arrow_assembly.AddPart(line_actor)
        #arrow_assembly.AddPart(label_actor)
        return arrow_assembly


class AxesWidget(AxesActor):
    def __init__(self,model):
        """
        the axes will be associated with the model that is given.
        """
        self.model = model
        length = np.sqrt(model.GetLength())
        Xaxis = self.make_arrow((1,0,0),'X', length, (1,0,0))
        Yaxis = self.make_arrow((0,1,0),'Y', length, (0,1,0))
        Zaxis = self.make_arrow((0,0,1),'Z', length, (0,0,1))
        self.AddPart(Xaxis)
        self.AddPart(Yaxis)
        self.AddPart(Zaxis)

        self.model_origin_trans = vtk.vtkTransform()
        self.model_origin_trans.SetMatrix(self.model.GetMatrix())
        self.model_origin_trans.Translate(self.model.GetOrigin())
        self.SetUserTransform(self.model_origin_trans)
        #self.SetOrientation(self.model.GetOrientation())
        
    def SetInteractor(self,iren):
        self.iren = iren
        self.iren.AddObserver("CharEvent", self.ProcessEvents)
        self.iren.AddObserver("MouseMoveEvent", self.ProcessEvents)
        
    def ProcessEvents(self,obj,event):
        key = self.iren.GetKeySym()
        if key==None:
            # Qt evnts
            try:
                key = self.iren._key
            except:
                pass
        
        #shift = self.iren.GetShiftKey()
        #ctrl = self.iren.GetControlKey()
        #print key
        if (key == 65) or (key == 'a'):
            self.ToggleVisibility()
 
        self.Update()
    
    def Update(self):
        self.model_origin_trans.SetMatrix(self.model.GetMatrix())
        self.model_origin_trans.Translate(self.model.GetOrigin())
        self.model_origin_trans.Update()

        
class drrPlaybackRep(JT.vtk.vtkPlaybackRepresentation):
    def __init__(self):
        #JT.vtk.vtkPlaybackRepresentation.__init__(self)
        prop = self.GetProperty()
        prop.SetColor(1,0,0)
        self.AddObserver("Play",self.Play)
    def Play(self,obj,event): 
        print "play\n"
    def Stop(self,obj,event): 
        print "stop\n"
    def ForwardOneFrame(self,obj,event): 
        print "forward one frame\n"
    def BackwardOneFrame(self,obj,event): 
        print "backward one frame\n"
    def JumpToBeginning(self,obj,event): 
        print "jump to beginning\n"
    def JumpToEnd(self,obj,event): 
        print "jump to end\n"

    
def test_drr():
    """
    create a drr render window
    """
    QtGui.QApplication.setStyle('cleanlooks')
    app = QtGui.QApplication(sys.argv) 
    mainWindow = QtGui.QMainWindow()
    frameWidget = QtGui.QFrame(mainWindow)

    mainWindow.setCentralWidget(frameWidget)
    drr = QTDRR(frameWidget)
     
    mainWindow.setFixedSize(512,512)
    drr.iren.setFixedSize(512,512)

    #canvas.iren.setFixedSize(512,512)
    image_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
    volume_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mhd'
    #volume_path = r'C:\Users\bryan\phd\data\Spine3\CT\Spine3.mhd'
    sta_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    cal_path = 'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim040-LAT.cal'
    
    drr.SetVolumeFileName(volume_path)
    drr.LoadCalFile(cal_path, 'external')
    drr.LoadStaFile(sta_path)
    #drr.AddAxesAtModelOrigin(drr.volume)
    axes = AxesWidget(drr.volume)
    drr.ren.AddActor(axes)
    #drr.AddSphereAtBodyCenter(drr.volume.GetPosition())
    #drr.volume.UseRayCastMapper()
    #drr.xrayCamera.SetCameraClippingRange(1,3000)
    
    background = JT.utils.genericImageReader(image_path)

    def OnChar(obj,event):
        #key = obj.GetKeySym()
        key = drr.iren._key
        #print key
        if key == 66:
            print "Set background image"
            drr.SetBackgroundImage(background)
        if key == 78:
            print "Set background image to none"
            drr.SetBackgroundImage(None)
        if key == 80: #p
            print "X rotation"
            for ang in range(-45,46,5):
                drr.volume.SetOrientation(ang,0,0)
                drr.Render()
            print "Y rotation"
            for ang in range(-45,46,5):
                drr.volume.SetOrientation(0,ang,0)
                drr.Render()
            print "Z rotation"
            for ang in range(-45,46,5):
                drr.volume.SetOrientation(0,0,ang)
                drr.Render()
            
            
    
    drr.iren.AddObserver("CharEvent", OnChar)

    drr.iren.Initialize()
    drr.Render()
    drr.iren.Start()
    
    mainWindow.show() 
    app.exec_()

def test_surf():
    """
    create a surface render window
    """
    
    QtGui.QApplication.setStyle('cleanlooks')
    app = QtGui.QApplication(sys.argv) 
    mainWindow = QtGui.QMainWindow()
    frameWidget = QtGui.QFrame(mainWindow)

    mainWindow.setCentralWidget(frameWidget)
    drr = QTDRR(frameWidget)
     
    mainWindow.setFixedSize(512,512)
    drr.iren.setFixedSize(512,512)

    #canvas.iren.setFixedSize(512,512)
    image_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
    volume_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mhd'
    sta_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    cal_path = 'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim040-LAT.cal'
    
    drr.AddAxesAtModelOrigin(drr.volume)
    surf = JT.SurfaceModel()
    surf.SetVolumeFileName(volume_path)
    surf.CreateIsoSurface(85,.5)
    surf.SetOrigin(drr.volume.GetOrigin())
    trans = vtk.vtkTransform()
    trans.SetMatrix(drr.volume.GetMatrix())
    
#    surf.SetUserTransform(drr.volume.GetUserTransform())
#    surf.SetOrientation(drr.volume.GetOrientation())
#    surf.SetPosition(drr.volume.GetPosition())
    surf.SetUserTransform(trans)
    drr.ren.AddActor(surf)
    drr.SetVolumeFileName(volume_path)
    drr.LoadCalFile(cal_path, 'external')
    drr.LoadStaFile(sta_path)

    background = JT.utils.genericImageReader(image_path)
    
    
    writer = vtk.vtkSTLWriter()
    writer.SetFileTypeToASCII()
    writer.SetInput(surf._surfaceModel)

    def OnChar(obj,event):
        #key = obj.GetKeySym()
        key = drr.iren._key
        print key
        if key == 66: #B
            drr.SetBackgroundImage(background)
        if key == 78: #N
            drr.SetBackgroundImage(None)
        if key == 49: #1
            drr.volume.SetVisibility(1)
            surf.SetVisibility(1)
            drr.Render() 
        if key == 50: #2
            drr.volume.SetVisibility(0)
            surf.SetVisibility(0)
            drr.Render() 
        if key == 51: #3
            drr.volume.SetVisibility(1)
            surf.SetVisibility(0)
            drr.Render() 
        if key == 52: #4
            drr.volume.SetVisibility(0)
            surf.SetVisibility(1)
            drr.Render() 
        if key == 83: #s
            filename = QtGui.QFileDialog.getSaveFileName()
            writer.SetFileName(str(filename))
            writer.Write()
            print "STL model saved!"
    
    def SyncProps(obj,event):
            drr.Render()
            if drr.volume == drr.style._PickedActor:
                surf.SetPosition(drr.volume.GetPosition())
                surf.SetOrientation(drr.volume.GetOrientation())
            elif surf == drr.style._PickedActor:
                drr.volume.SetPosition(surf.GetPosition())
                drr.volume.SetOrientation(surf.GetOrientation())
            
    drr.iren.AddObserver("KeyPressEvent",SyncProps)
    drr.iren.AddObserver("MouseMoveEvent",SyncProps)
    
    drr.iren.AddObserver("CharEvent", OnChar)
    
    drr.iren.Initialize()
    drr.Render()
    drr.iren.Start()
    
    mainWindow.show() 
    app.exec_()

def test_free_view():
# Setup a free view that will display the 3d volume from any angle
    QtGui.QApplication.setStyle('cleanlooks')
    app = QtGui.QApplication(sys.argv) 
    mainWindow = QtGui.QMainWindow()
    
    freeViewFrame = QtGui.QFrame()
    freeViewFrame.setMinimumSize(400,400)
    freeViewIren = QVTKRenderWindowInteractor(freeViewFrame)
    freeViewRen = JT.vtk.vtkRenderer()
    freeViewRenWin = freeViewIren.GetRenderWindow()
    freeViewRenWin.AddRenderer(freeViewRen)
    
    mainWindow.setCentralWidget(freeViewFrame)
    
    cone = vtk.vtkConeSource()
    cone.SetResolution(30)

    coneMapper = vtk.vtkPolyDataMapper()
    coneMapper.SetInput(cone.GetOutput())

    coneActor = vtk.vtkActor()
    coneActor.SetMapper(coneMapper)
    
    freeViewRen.AddActor(coneActor)
    freeViewIren
    freeViewRenWin.Render()
    freeViewIren.Initialize()
    freeViewIren.Start()
    
    mainWindow.show() 
    app.exec_()


def test_stl():
    """
    create a drr render window
    """
    QtGui.QApplication.setStyle('cleanlooks')
    app = QtGui.QApplication(sys.argv) 
    mainWindow = QtGui.QMainWindow()
    frameWidget = QtGui.QFrame(mainWindow)

    mainWindow.setCentralWidget(frameWidget)
    drr = QTDRR(frameWidget)
     
    mainWindow.setFixedSize(512,512)
    drr.iren.setFixedSize(512,512)
    #style = vtk.vtkInteractorStyleTrackballActor()
    #drr.iren.SetInteractorStyle(style)

    volume_path = r'C:\Users\bryan\phd\data\Models\Synthes CAD Files\09_820_025_2ascii.stl'
    volume_path2 = r'C:\Users\bryan\phd\data\Models\Synthes CAD Files\09_820_025_1ascii.stl'
    image_path = r'C:\Users\bryan\phd\data\Spine1\Spine1_1\Movies\crop_512\Spine1_FE1_0001.tif'
    
    cal_path = r'C:\Users\bryan\phd\data\Spine1\Spine1_1\Movies\crop_512\Calibration_cal_m.cal'
    
    #drr.SetVolumeFileName(volume_path)
    drr.LoadCalFile(cal_path, 'internal')
    #drr.LoadStaFile(sta_path)
    #drr.xrayCamera.SetCameraClippingRange(1,3000)
    
    #background = JT.utils.genericImageReader(image_path)
    #drr.SetBackgroundImage(background)
    
    drr.LoadSTL(volume_path)
    vol2 = JT.SurfaceModel()
    vol2.LoadSTL(volume_path2)
    drr.ren.AddActor(vol2)
    
    w2if = vtk.vtkWindowToImageFilter()
    w2if.SetInput(drr.renWin)
    save_name = r'C:\Users\bryan\phd\data\Analysis\ch2_implant_rot'
    writer = vtk.vtkPNGWriter()
    writer.SetInput(w2if.GetOutput())
    #drr.volume.SetPosition(400,0,0)
    def OnChar(obj,event):
        #key = obj.GetKeySym()
        key = drr.iren._key
        print key
        if key == 66: #B
            # loop through a series of rotations and save a projection image for each one.
            for i,ang in enumerate(range(-5,6,1)):
                drr.volume.SetOrientation(ang,0,0)
                drr.Render()
                w2if.Modified()
                writer.SetFileName(save_name+'X_'+str(i)+'.png')
                writer.Write()
            for i,ang in enumerate(range(-5,6,1)):
                drr.volume.SetOrientation(0,ang,0)
                drr.Render()
                w2if.Modified()
                writer.SetFileName(save_name+'Y_'+str(i)+'.png')
                writer.Write()
            for i,ang in enumerate(range(-5,6,1)):
                drr.volume.SetOrientation(0,0,ang)
                drr.Render()
                w2if.Modified()
                writer.SetFileName(save_name+'Z_'+str(i)+'.png')
                writer.Write()     
#            print "initial position %s" % str(drr.xrayCamera.GetPosition())
#            drr.ren.ResetCamera()
#            print "reset position %s" % str(drr.xrayCamera.GetPosition())
#            print "actor posiiton %s" % str(drr.volume.GetPosition())

    drr.iren.AddObserver("CharEvent", OnChar)

    
    drr.Render()
    drr.iren.Initialize()
    drr.iren.Start()
    
    mainWindow.show() 
    app.exec_()

def test_nongui_drr():
    """
    create a drr render window
    """
    drr = DRR()
     
    drr.renWin.SetSize(512,512)

    #canvas.iren.setFixedSize(512,512)
    image_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/fluoro/ushortim080-LAT.mhd'
    #volume_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.mhd'
    volume_path = r'C:\Users\bryan\phd\data\CT\Spine3.mhd'
    sta_path = 'C:/Users/bryan/bryan-code/2D3D/vert1/CT/ucharCTv1r1.sta'
    cal_path = 'C:/Users/bryan/bryan-code/2D3D/vert2/fluoro/ushortim040-LAT.cal'
    
    drr.SetVolumeFileName(volume_path)
    drr.LoadCalFile(cal_path, 'external')
    drr.LoadStaFile(sta_path)

    #drr.xrayCamera.SetCameraClippingRange(1,3000)
    
    background = JT.utils.genericImageReader(image_path)

    drr.iren.Initialize()
    drr.Render()
    drr.iren.Start()

if __name__ == '__main__':
    test_drr()
#    test_stl()